import os
import sys
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib
from tqdm import tqdm
import cv2
from torchvision import transforms
from skimage.color import rgb2lab, lab2rgb

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 导入模型
from models.colorization_model import ColorizationGenerator

# 设置中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']  # 优先使用的中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
font_path = '/Users/liuyuzhuo/Library/Fonts/SimHei.ttf'  # 请确保此路径正确

# 尝试设置中文字体，如果文件存在
if os.path.exists(font_path):
    from matplotlib.font_manager import FontProperties
    font_prop = FontProperties(fname=font_path)
    plt.rcParams['font.family'] = font_prop.get_name()
else:
    print("警告：找不到中文字体文件，可能导致中文显示为乱码")

def load_model(model_path, device=None):
    """加载训练好的模型"""
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 加载检查点
    print(f"加载模型文件: {model_path}")
    checkpoint = torch.load(model_path, map_location=device)
    
    # 创建模型
    model = ColorizationGenerator()
    
    # 加载模型状态字典
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()
    
    # 打印模型信息
    print(f"模型加载成功，epoch: {checkpoint.get('epoch', 'unknown')}")
    
    return model, device

def preprocess_image(image_path, img_size=256):
    """预处理图像，确保与训练时的预处理一致"""
    # 读取图像
    img = Image.open(image_path).convert('RGB')
    
    # 调整大小
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
    ])
    
    img_tensor = transform(img)
    
    # 转换为LAB颜色空间
    img_lab = rgb2lab(img_tensor.permute(1, 2, 0).numpy())
    
    # 分离L通道和ab通道
    L = img_lab[:, :, 0]
    ab = img_lab[:, :, 1:]
    
    # 归一化
    L = L / 50.0 - 1.0  # 归一化到[-1, 1]
    ab = ab / 110.0  # 归一化到[-1, 1]
    
    # 转换为张量
    L_tensor = torch.from_numpy(L).unsqueeze(0).float()  # [1, H, W]
    ab_tensor = torch.from_numpy(ab).permute(2, 0, 1).float()  # [2, H, W]
    
    return L_tensor, ab_tensor, img_tensor

def colorize_image(model, L_tensor, device):
    """使用模型为灰度图像上色"""
    with torch.no_grad():
        # 将图像移到设备
        L_tensor = L_tensor.to(device)
        
        # 确保输入有批次维度 (B,C,H,W)
        if L_tensor.dim() == 3:  # 如果是(C,H,W)
            L_tensor = L_tensor.unsqueeze(0)  # 变成(1,C,H,W)
        
        # 生成着色结果
        ab_fake = model(L_tensor)
        
    return ab_fake.cpu()  # 返回CPU张量

def lab_to_rgb(L, ab):
    """
    将L通道和ab通道转换为RGB图像
    
    Args:
        L: 形状为[H, W]或[1, H, W]的L通道张量，值范围为[-1, 1]
        ab: 形状为[2, H, W]的ab通道张量，值范围为[-1, 1]
    
    Returns:
        RGB图像，形状为[H, W, 3]，值范围为[0, 1]
    """
    # 确保L和ab在CPU上
    L = L.cpu()
    ab = ab.cpu()
    
    # 确保L是2D张量
    if L.dim() > 2:
        L_np = L.squeeze().numpy()
    else:
        L_np = L.numpy()
    
    # 将L通道从[-1, 1]转换为[0, 100]
    L_np = (L_np + 1.0) * 50.0
    L_np = np.clip(L_np, 0, 100)
    
    # 将ab通道从[-1, 1]转换为[-128, 127]
    ab_np = ab.numpy() * 110.0
    ab_np = np.clip(ab_np, -128, 127)
    
    # 创建LAB图像
    lab_image = np.zeros((L_np.shape[0], L_np.shape[1], 3))
    lab_image[:, :, 0] = L_np
    
    # 确保ab_np的维度正确
    # 如果ab_np形状是[2, H, W]，则需要将其重新排列为[H, W, 2]
    if ab_np.shape[0] == 2:
        lab_image[:, :, 1:] = np.transpose(ab_np, (1, 2, 0))
    else:
        lab_image[:, :, 1:] = ab_np.reshape(L_np.shape[0], L_np.shape[1], 2)
    
    # 转换为RGB
    with np.errstate(invalid='ignore'):  # 忽略无效值警告
        rgb_image = lab2rgb(lab_image)
    
    return rgb_image

def create_grid_comparison(images_data, output_path, model_type):
    """创建3x3网格图像比较
    
    Args:
        images_data: 包含图像处理结果的列表，每项是一个字典包含L、ab_real、ab_fake
        output_path: 输出路径
        model_type: 模型类型，用于设置标题
    """
    # 创建一个3x3的图表，设置更大的分辨率
    fig, axes = plt.subplots(3, 3, figsize=(18, 18), dpi=150)
    
    # 确保我们有3个样本
    if len(images_data) < 3:
        # 如果样本少于3个，重复最后一个样本填充
        while len(images_data) < 3:
            images_data.append(images_data[-1])
    elif len(images_data) > 3:
        # 如果样本多于3个，只使用前3个
        images_data = images_data[:3]
    
    # 设置列标题，增大字体
    fontsize = 14
    axes[0, 0].set_title("灰度图像", fontproperties=font_prop if 'font_prop' in globals() else None, fontsize=fontsize)
    axes[0, 1].set_title("真实彩色", fontproperties=font_prop if 'font_prop' in globals() else None, fontsize=fontsize)
    axes[0, 2].set_title("生成彩色", fontproperties=font_prop if 'font_prop' in globals() else None, fontsize=fontsize)
    
    # 显示每个样本的不同变体
    for i, data in enumerate(images_data):
        # 确保所有数据都在CPU上
        L = data['L'].cpu()
        ab_real = data['ab_real'].cpu()
        ab_fake = data['ab_fake'].cpu()
        filename = data.get('filename', f'样本{i+1}')
        
        # 转换为RGB显示格式
        gray_img = (L.squeeze().numpy() + 1.0) / 2.0  # 从[-1, 1]转换为[0, 1]
        real_rgb = lab_to_rgb(L, ab_real)
        fake_rgb = lab_to_rgb(L, ab_fake)
        
        # 显示灰度图像
        axes[i, 0].imshow(gray_img, cmap='gray')
        axes[i, 0].axis('off')
        
        # 显示真实彩色图像
        axes[i, 1].imshow(real_rgb)
        axes[i, 1].axis('off')
        
        # 显示生成的彩色图像
        axes[i, 2].imshow(fake_rgb)
        axes[i, 2].axis('off')
        
        # 可以在每行左侧添加行标题
        if i == 0:
            axes[i, 0].set_ylabel("样本1", fontproperties=font_prop if 'font_prop' in globals() else None, fontsize=fontsize)
        elif i == 1:
            axes[i, 0].set_ylabel("样本2", fontproperties=font_prop if 'font_prop' in globals() else None, fontsize=fontsize)
        elif i == 2:
            axes[i, 0].set_ylabel("样本3", fontproperties=font_prop if 'font_prop' in globals() else None, fontsize=fontsize)
    
    # 设置整体标题，增大字体
    title = f"{model_type}模型着色效果"
    if 'font_prop' in globals():
        plt.suptitle(title, fontproperties=font_prop, fontsize=24)
    else:
        plt.suptitle(title, fontsize=24)
    
    # 调整布局，确保不会裁剪
    plt.tight_layout()
    plt.subplots_adjust(top=0.95, hspace=0.1, wspace=0.1)
    
    # 保存图像，设置较高的DPI
    plt.savefig(output_path, dpi=200, bbox_inches='tight')
    plt.close()

def create_html_summary(flower_output, face_output, output_dir):
    """创建HTML摘要页面"""
    html_path = os.path.join(output_dir, "index.html")
    with open(html_path, 'w', encoding='utf-8') as f:
        f.write("<!DOCTYPE html>\n")
        f.write("<html>\n<head>\n")
        f.write("<meta charset='utf-8'>\n")
        f.write("<title>模型着色效果比较</title>\n")
        f.write("<style>body{font-family:'Microsoft YaHei',Arial,sans-serif;margin:20px;text-align:center;}h1,h2{color:#333;}img{max-width:100%;margin-bottom:20px;}</style>\n")
        f.write("</head>\n<body>\n")
        f.write("<h1>模型着色效果比较</h1>\n")
        
        # 显示花卉图像
        if flower_output:
            f.write("<h2>花卉模型着色效果</h2>\n")
            f.write(f"<div style='margin-bottom:30px;'>\n")
            f.write(f"<img src='{os.path.basename(flower_output)}'>\n")
            f.write("</div>\n")
        
        # 显示人脸图像
        if face_output:
            f.write("<h2>人脸模型着色效果</h2>\n")
            f.write(f"<div style='margin-bottom:30px;'>\n")
            f.write(f"<img src='{os.path.basename(face_output)}'>\n")
            f.write("</div>\n")
        
        f.write("</body>\n</html>")
    
    print(f"已创建HTML摘要: {html_path}")

def main():
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 设置路径
    flower_samples_dir = "data/flower/samples"
    face_samples_dir = "data/mini_face/samples"
    output_dir = "output/cn_samples"
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 检查模型文件
    flower_model_path = "models/weights/flower.pth"
    face_model_path = "models/weights/face.pth"
    
    if not os.path.exists(flower_model_path):
        print(f"错误: 找不到花卉模型文件 {flower_model_path}")
        return
    
    if not os.path.exists(face_model_path):
        print(f"错误: 找不到人脸模型文件 {face_model_path}")
        return
    
    # 加载模型
    print("正在加载模型...")
    flower_model, _ = load_model(flower_model_path, device)
    face_model, _ = load_model(face_model_path, device)
    print("模型加载完成!")
    
    # 存储处理结果
    flower_data = []
    face_data = []
    
    # 处理花卉样本
    print("处理花卉样本...")
    if os.path.exists(flower_samples_dir):
        # 获取所有JPG图像
        flower_files = [f for f in os.listdir(flower_samples_dir) if f.endswith('.jpg') or f.endswith('.png')]
        
        # 只处理最多3个样本
        sample_count = min(3, len(flower_files))
        for file in tqdm(flower_files[:sample_count], desc="花卉样本"):
            image_path = os.path.join(flower_samples_dir, file)
            
            # 预处理图像
            L_tensor, ab_real, _ = preprocess_image(image_path)
            
            # 生成着色结果
            ab_fake = colorize_image(flower_model, L_tensor, device)
            
            # 存储结果
            flower_data.append({
                'L': L_tensor,
                'ab_real': ab_real,
                'ab_fake': ab_fake,
                'filename': file
            })
    else:
        print(f"警告: 找不到花卉样本目录 {flower_samples_dir}")
    
    # 处理人脸样本
    print("处理人脸样本...")
    if os.path.exists(face_samples_dir):
        # 获取所有JPG图像
        face_files = [f for f in os.listdir(face_samples_dir) if f.endswith('.jpg') or f.endswith('.png')]
        
        # 只处理最多3个样本
        sample_count = min(3, len(face_files))
        for file in tqdm(face_files[:sample_count], desc="人脸样本"):
            image_path = os.path.join(face_samples_dir, file)
            
            # 预处理图像
            L_tensor, ab_real, _ = preprocess_image(image_path)
            
            # 生成着色结果
            ab_fake = colorize_image(face_model, L_tensor, device)
            
            # 存储结果
            face_data.append({
                'L': L_tensor,
                'ab_real': ab_real,
                'ab_fake': ab_fake,
                'filename': file
            })
    else:
        print(f"警告: 找不到人脸样本目录 {face_samples_dir}")
    
    # 创建并保存网格比较图
    flower_output = ""
    face_output = ""
    
    if flower_data:
        flower_output = os.path.join(output_dir, "flower_comparison.png")
        create_grid_comparison(flower_data, flower_output, "花卉")
        print(f"已创建花卉网格图像: {flower_output}")
    
    if face_data:
        face_output = os.path.join(output_dir, "face_comparison.png")
        create_grid_comparison(face_data, face_output, "人脸")
        print(f"已创建人脸网格图像: {face_output}")
    
    # 创建HTML摘要
    create_html_summary(flower_output, face_output, output_dir)
    
    print(f"处理完成! 生成了花卉和人脸模型的着色效果比较图")
    print(f"结果保存在: {os.path.abspath(output_dir)}")

if __name__ == "__main__":
    main() 